'''
Author: daniel
Date: 2023-12-17 20:27:24
LastEditTime: 2023-12-17 20:35:39
LastEditors: daniel
Description: 
FilePath: /topictrack-bee/tools/det_acc.py
have a nice day
'''
import numpy as np




def iou(box1, box2):
    """
    计算两个框之间的交并比
    box1: [xmin, ymin, xmax, ymax]
    box2: [xmin, ymin, xmax, ymax]
    """
    x1 = max(box1[3], box2[3])
    y1 = max(box1[4], box2[4])
    x2 = min(box1[5], box2[5])
    y2 = min(box1[6], box2[6])
    inter_area = max(0, x2 - x1) * max(0, y2 - y1)

    box1_area = (box1[5] - box1[3]) * (box1[6] - box1[4])
    box2_area = (box2[5] - box2[3]) * (box2[6] - box2[4])
    union_area = box1_area + box2_area - inter_area

    return inter_area / union_area


def compute_ap(gt_boxes, pred_boxes, iou_threshold=0.5):
    """
    计算 AP
    gt_boxes: 真实框的列表, [num_gt_boxes, 4]，每一行是一个框的坐标
    pred_boxes: 预测框的列表，[num_pred_boxes, 4]，每一行是一个框的坐标
    iou_threshold: IOU 阈值，预测与真实框之间的 IOU 大于此值则被视为正确的预测
    """
    num_gt_boxes = len(gt_boxes)
    num_pred_boxes = len(pred_boxes)

    # 如果没有预测框，则返回AP为0
    if num_pred_boxes == 0:
        return 0.0

    # 按置信度对预测框排序
    pred_boxes = np.array(pred_boxes)

    sorted_pred_boxes_inds = np.argsort(-pred_boxes[:, 2])
    sorted_pred_boxes = pred_boxes[sorted_pred_boxes_inds]

    # 记录每个预测框是否正确
    true_positives = np.zeros(num_pred_boxes)
    false_positives = np.zeros(num_pred_boxes)

    for i in range(num_pred_boxes):
        pred_box = sorted_pred_boxes[i]

        # 与真实框计算交并比
        ious = np.array([iou(pred_box, gt_box) for gt_box in gt_boxes])

        max_iou_index = np.argmax(ious)
        max_iou = ious[max_iou_index]

        # 如果 IOU 高于阈值，则认为预测正确

        if max_iou >= iou_threshold:
            # 真实框已被匹配，预测框为重复预测
            if max_iou_index >= len(true_positives):
                continue
            if true_positives[max_iou_index] == 1:
                false_positives[i] = 1
            # 否则标记为正确预测
            else:
                true_positives[max_iou_index] = 1
        # 如果 IOU 低于阈值，则认为预测错误
        else:
            false_positives[i] = 1

    # 计算 AP
    cum_true_positives = np.cumsum(true_positives)
    cum_false_positives = np.cumsum(false_positives)
    precision = cum_true_positives / (cum_true_positives + cum_false_positives + 1e-10)
    recall = cum_true_positives / num_gt_boxes
    ap = 0.0
    for i in range(11):
        t = i / 10.0
        index = np.where(recall >= t)
        if len(index[0]) == 0:
            p = 0
        else:
            p = np.max(precision[index])
        ap += p / 11.0
    return ap


def SaveInfo(info, result_path):
    record = ','.join(list(map(str, info)))

    with open(result_path, "a+") as f:
        f.write(record+"\n")





def cal_det_ap(gt_path, pre_path, iou_th):
    gt = []
    for line in open(gt_path):
        data = line.split(',')
        if int(float(data[2])) < 0:
            data[2] = 0
        if int(float(data[3])) < 0:
            data[3] = 0
        info = [int(data[0]), 0, 1, int(float(data[2])), int(float(data[3])), int(float(data[2])) + int(float(data[4])), int(float(data[3])) + int(float(data[5]))]
        gt.append(info)
    gt = np.array(gt)
    predicted = []
    for line in open(pre_path):
        data = line.split(',')
        if int(float(data[2])) < 0:
            data[2] = 0
        if int(float(data[3])) < 0:
            data[3] = 0
        info = [int(data[0]), 0, 1, int(float(data[2])), int(float(data[3])), int(float(data[2])) + int(float(data[4])), int(float(data[3])) + int(float(data[5]))]
        info1 = []
        predicted.append(info)
    predicted = np.array(predicted)

    ap_all = []
    for i in range(1,gt[-1][0]+1):
        # print('id ',i)
        gt_one = [row for row in gt if row[0] == i]
        pre_one = [row for row in predicted if row[0] == i]
        ap = compute_ap(gt_one,pre_one,iou_th)
        ap_all.append(ap)

    return sum(ap_all)/gt[-1][0]

